import random
from threading import local
import time
import numpy as np
from isaacgym import gymutil
from isaacgym import gymapi
#from isaacgym import gymtorch
from math import sqrt
import math
from omegaconf import read_write
from sympy import false, re
import torch
import cv2
import gym
from gym import spaces

from draw import *

from pcgworker.PCGWorker import *

from stable_baselines3.common.env_checker import check_env

# https://stable-baselines3.readthedocs.io/en/master/guide/custom_env.html

N_DISCRETE_ACTIONS = 6
N_CHANNELS = 3
HEIGHT = 84
WIDTH = 84

class CustomEnv():
    """Custom Environment that follows gym interface"""
    metadata = {'render.modes': ['human', 'rgb_array']}

    # Set this in SOME subclasses
    metadata = {"render.modes": []}
    reward_range = (-float("inf"), float("inf"))
    spec = None

    # Set these in ALL subclasses
    action_space = None
    observation_space = None

    def __init__(self, headless_: bool = True, seed_pth: str = "blank_wave.json", render_indicator: bool = True):
        super(CustomEnv, self).__init__()
        # Define action and observation space
        self.action_space = spaces.Discrete(N_DISCRETE_ACTIONS)
        # self.observation_space = spaces.Box(low=0, high=255, shape=(HEIGHT, WIDTH, N_CHANNELS), dtype=np.uint8)
        self.observation_space = spaces.Box(low=0, high=255, shape=(N_CHANNELS, HEIGHT, WIDTH), dtype=np.uint8)

        torch.manual_seed(3407)

        '''
        simulatiom parameters
        '''
        self.num_matrix_envs = 1
        self.headless = headless_
        self.spd = 10
        max_actors_num_per_env = 9*9*4+1

        self.render_indicator = render_indicator

        self.initial_height = 2.05

        # get default set of parameters
        sim_params = gymapi.SimParams()
        args = gymutil.parse_arguments()
        sim_params.physx.solver_type = 1
        sim_params.physx.num_position_iterations = 2
        sim_params.physx.num_velocity_iterations = 1
        sim_params.physx.num_threads = args.num_threads
        sim_params.physx.use_gpu = args.use_gpu
        sim_params.use_gpu_pipeline = False #args.use_gpu_pipeline
        sim_params.substeps = 2
        sim_params.dt = 1.0 / 30.0
        # setting up the Z-up axis
        sim_params.up_axis = gymapi.UP_AXIS_Z
        sim_params.gravity = gymapi.Vec3(0.0, 0.0, -9.8)

        # configure the ground plane
        plane_params = gymapi.PlaneParams()
        plane_params.normal = gymapi.Vec3(0, 0, 1) # z-up!
        plane_params.distance = 1
        plane_params.static_friction = 1
        plane_params.dynamic_friction = 1
        plane_params.restitution = 0.1

        # set up the env grid
        num_envs = 1
        spacing = 20
        env_lower = gymapi.Vec3(-spacing, -spacing, -spacing)
        env_upper = gymapi.Vec3(spacing, spacing, spacing)

        # Attractor setup
        attractor_properties = gymapi.AttractorProperties()
        attractor_properties.stiffness = 5e10
        attractor_properties.damping = 0
        attractor_properties.axes = gymapi.AXIS_ROTATION
        attractor_pose = gymapi.Transform()
        attractor_pose.p = gymapi.Vec3(0, 0.0, 0.0)
        attractor_pose.r = gymapi.Quat(0,1,0,1)

        # create procedural asset
        asset_options = gymapi.AssetOptions()
        asset_options.density = 1.0
        asset_options.fix_base_link = False
        asset_options.linear_damping = 0.1
        asset_options1 = gymapi.AssetOptions()
        asset_options1.density = 0.02
        asset_options1.fix_base_link = True
        asset_options1.linear_damping = 0.1

        color_red = gymapi.Vec3(1,0,0)
        color_green = gymapi.Vec3(0,1,0)

        # load PCG asset
        asset_root = "./assets"
        pt_gray_cube = "PCG/gray_cube.urdf"
        pt_blue_cube = "PCG/blue_cube.urdf"
        pt_red_cube = "PCG/red_cube.urdf"
        pt_white_cube = "PCG/white_cube.urdf"
        pt_yellow_cube = "PCG/yellow_cube.urdf"
        pt_blue_ramp = "PCG/blue_ramp.urdf"
        pt_red_ramp = "PCG/red_ramp.urdf"
        pt_white_ramp = "PCG/white_ramp.urdf"
        pt_yellow_ramp = "PCG/yellow_ramp.urdf"
        pt_blue_corner = "PCG/blue_corner.urdf"
        pt_red_corner = "PCG/red_corner.urdf"
        pt_white_corner = "PCG/white_corner.urdf"
        pt_yellow_corner = "PCG/yellow_corner.urdf"


        # ramp and corner attitutde
        ramp_rot0 = gymapi.Quat.from_axis_angle(gymapi.Vec3(1, 0, 0), 0.5*math.pi) * gymapi.Quat.from_axis_angle(gymapi.Vec3(0, 1, 0), 0*math.pi)
        ramp_rot1 = gymapi.Quat.from_axis_angle(gymapi.Vec3(1, 0, 0), 0.5*math.pi) * gymapi.Quat.from_axis_angle(gymapi.Vec3(0, 1, 0), 0.5*math.pi)
        ramp_rot2 = gymapi.Quat.from_axis_angle(gymapi.Vec3(1, 0, 0), 0.5*math.pi) * gymapi.Quat.from_axis_angle(gymapi.Vec3(0, 1, 0), 1*math.pi)
        ramp_rot3 = gymapi.Quat.from_axis_angle(gymapi.Vec3(1, 0, 0), 0.5*math.pi) * gymapi.Quat.from_axis_angle(gymapi.Vec3(0, 1, 0), 1.5*math.pi)
        corner_rot0 = gymapi.Quat.from_axis_angle(gymapi.Vec3(0, 1, 0), 0.5*math.pi) * gymapi.Quat.from_axis_angle(gymapi.Vec3(1, 0, 0), 0*math.pi)
        corner_rot1 = gymapi.Quat.from_axis_angle(gymapi.Vec3(0, 1, 0), 0.5*math.pi) * gymapi.Quat.from_axis_angle(gymapi.Vec3(1, 0, 0), 0.5*math.pi)
        corner_rot2 = gymapi.Quat.from_axis_angle(gymapi.Vec3(0, 1, 0), 0.5*math.pi) * gymapi.Quat.from_axis_angle(gymapi.Vec3(1, 0, 0), 1*math.pi)
        corner_rot3 = gymapi.Quat.from_axis_angle(gymapi.Vec3(0, 1, 0), 0.5*math.pi) * gymapi.Quat.from_axis_angle(gymapi.Vec3(1, 0, 0), 1.5*math.pi)

        # birick size
        interval_xy = 2
        interval_z = 2*0.7

        # camera sensor properties
        camera_properties = gymapi.CameraProperties()
        camera_properties.width = WIDTH
        camera_properties.height = HEIGHT
        camera_properties.enable_tensors = True

        # rigid shape material properties
        shape_props = [gymapi.RigidShapeProperties()]
        shape_props[0].friction = -1
        shape_props[0].rolling_friction = 1
        shape_props[0].torsion_friction = 1
        shape_props[0].compliance = 0
        shape_props[0].restitution = 0
        shape_props1 = [gymapi.RigidShapeProperties()]
        shape_props1[0].friction = 0.2
        shape_props1[0].rolling_friction = -1
        shape_props1[0].torsion_friction = -1
        shape_props1[0].compliance = 0
        shape_props1[0].restitution = 0
        shape_props2 = [gymapi.RigidShapeProperties()]
        shape_props2[0].friction = 1
        shape_props2[0].rolling_friction = -1
        shape_props2[0].torsion_friction = -1
        shape_props2[0].compliance = 0
        shape_props2[0].restitution = 0

        # single world camera image
        self.camera_sensor_image = np.zeros((HEIGHT,WIDTH, 3), np.uint8)
        # human control
        self.interaction = 0

        # proximity threshold for reward calculation
        self.proximity_threshold = 1.5
        self.max_steps = 1000

        self.show_action = False

        '''
        generate WFC maps
        '''
        self.PCGWorker_ = PCGWorker(9,9)

        seed = self.PCGWorker_.from_file(filename = seed_pth)

        self.height_map = self.PCGWorker_.render(seed, output=True, border=False)

        # seed_ = PCGWorker_.from_file(filename = "blank_wave.json")
        # seed = PCGWorker_.mutate(seed_,81)
        # PCGWorker_.to_file(seed,filename = "seed.json")

        # seed = PCGWorker_.generate()
        self.connectivity_map , self.connectivity_map_visual = self.PCGWorker_.connectivity_analysis(seed,visualize_ = False)

        self.seeds = [seed]
        for i in range(0,self.num_matrix_envs-1):
            seed_ = self.PCGWorker_.mutate(self.seeds[-1],81)
            self.seeds.append(seed_)
        building_block_pointer = [0,0,0,0,0,0,0,0,0,0,0,0,0]    # 0-4,brick; 5-8,corners; 9-12 ,ramps;

        self.gym = gymapi.acquire_gym()

        self.sim = self.gym.create_sim(args.compute_device_id, args.graphics_device_id, gymapi.SIM_PHYSX, sim_params)

        print("args.graphics_device_id : ",args.graphics_device_id)
        print("args.compute_device_id : ",args.compute_device_id)


        # create the ground plane
        self.gym.add_ground(self.sim, plane_params)

        capsule_asset = self.gym.create_capsule(self.sim, 1, 1, asset_options)
        # capsule_asset = self.gym.create_box(self.sim, 1, 1, 1, asset_options)

        asset_gray_cube = self.gym.load_asset(self.sim, asset_root, pt_gray_cube,asset_options1)
        asset_blue_cube = self.gym.load_asset(self.sim, asset_root, pt_blue_cube,asset_options1)
        asset_red_cube = self.gym.load_asset(self.sim, asset_root, pt_red_cube,asset_options1)
        asset_white_cube = self.gym.load_asset(self.sim, asset_root, pt_white_cube,asset_options1)
        asset_yellow_cube = self.gym.load_asset(self.sim, asset_root, pt_yellow_cube,asset_options1)
        asset_blue_ramp = self.gym.load_asset(self.sim, asset_root, pt_blue_ramp,asset_options1)
        asset_red_ramp = self.gym.load_asset(self.sim, asset_root, pt_red_ramp,asset_options1)
        asset_white_ramp = self.gym.load_asset(self.sim, asset_root, pt_white_ramp,asset_options1)
        asset_yellow_ramp = self.gym.load_asset(self.sim, asset_root, pt_yellow_ramp,asset_options1)
        asset_blue_corner = self.gym.load_asset(self.sim, asset_root, pt_blue_corner,asset_options1)
        asset_red_corner = self.gym.load_asset(self.sim, asset_root, pt_red_corner,asset_options1)
        asset_white_corner = self.gym.load_asset(self.sim, asset_root, pt_white_corner,asset_options1)
        asset_yellow_corner = self.gym.load_asset(self.sim, asset_root, pt_yellow_corner,asset_options1)


        box_asset = self.gym.create_box(self.sim, 1, 1, 1, asset_options)

        if not self.headless:
            cam_props = gymapi.CameraProperties()
            self.viewer = self.gym.create_viewer(self.sim, cam_props)

        self.envs = []
        self.camera_handles = []
        self.actor_handles = []
        self.food_handles = []
        self.attractor_handles = []
        self.force_sensors = []

        for n in range(self.num_matrix_envs):
            # create environments
            env_ = self.gym.create_env(self.sim, env_lower, env_upper, 4)
            self.envs.append(env_)

            n_actors = 0
            # create actors based on WFC seed
            for i in range(0,9):
                for j in range(0,9):
                    tile_ = self.seeds[n].wave_oriented[j*9+i][0][0]
                    rot = self.seeds[n].wave_oriented[j*9+i][0][1]
                    pose_ij = gymapi.Transform()
                    pose_ij.p = gymapi.Vec3(i*interval_xy, j*interval_xy, 0*interval_xy)
                    pose_ij.r = gymapi.Quat(0,0,0,1)
                    actor = self.gym.create_actor(env_, asset_gray_cube, pose_ij, "gray_cube" + str(i) + str(j), n, 0)
                    # cubes
                    if tile_<=5:
                        if tile_==1:
                            pass
                        elif tile_==2:
                            pose_ij.p = gymapi.Vec3(i*interval_xy, j*interval_xy, 1*interval_z)
                            actor = self.gym.create_actor(env_, asset_blue_cube, pose_ij, "blue_cube" + str(i) + str(j), n, 0)
                        elif tile_==3:
                            pose_ij.p = gymapi.Vec3(i*interval_xy, j*interval_xy, 1*interval_z)
                            actor = self.gym.create_actor(env_, asset_blue_cube, pose_ij, "blue_cube" + str(i) + str(j), n, 0)
                            pose_ij.p = gymapi.Vec3(i*interval_xy, j*interval_xy, 2*interval_z)
                            actor = self.gym.create_actor(env_, asset_yellow_cube, pose_ij, "yellow_cube" + str(i) + str(j), n, 0)
                        elif tile_==4:
                            pose_ij.p = gymapi.Vec3(i*interval_xy, j*interval_xy, 1*interval_z)
                            actor = self.gym.create_actor(env_, asset_blue_cube, pose_ij, "blue_cube" + str(i) + str(j), n, 0)
                            pose_ij.p = gymapi.Vec3(i*interval_xy, j*interval_xy, 2*interval_z)
                            actor = self.gym.create_actor(env_, asset_yellow_cube, pose_ij, "yellow_cube" + str(i) + str(j), n, 0)
                            pose_ij.p = gymapi.Vec3(i*interval_xy, j*interval_xy, 3*interval_z)
                            actor = self.gym.create_actor(env_, asset_red_cube, pose_ij, "red_cube" + str(i) + str(j), n, 0)
                        elif tile_==5:
                            pose_ij.p = gymapi.Vec3(i*interval_xy, j*interval_xy, 1*interval_z)
                            actor = self.gym.create_actor(env_, asset_blue_cube, pose_ij, "blue_cube" + str(i) + str(j), n, 0)
                            pose_ij.p = gymapi.Vec3(i*interval_xy, j*interval_xy, 2*interval_z)
                            actor = self.gym.create_actor(env_, asset_yellow_cube, pose_ij, "yellow_cube" + str(i) + str(j), n, 0)
                            pose_ij.p = gymapi.Vec3(i*interval_xy, j*interval_xy, 3*interval_z)
                            actor = self.gym.create_actor(env_, asset_red_cube, pose_ij, "red_cube" + str(i) + str(j), n, 0)
                            pose_ij.p = gymapi.Vec3(i*interval_xy, j*interval_xy, 4*interval_z)
                            actor = self.gym.create_actor(env_, asset_white_cube, pose_ij, "white_cube" + str(i) + str(j), n, 0)
                    # corners
                    elif tile_==6:
                        if rot==0:
                            pose_ij.p = gymapi.Vec3(i*interval_xy, j*interval_xy, 1*interval_z)
                            pose_ij.r = corner_rot2
                            actor = self.gym.create_actor(env_, asset_blue_corner, pose_ij, "blue_corner" + str(i) + str(j), n, 0)
                        elif rot==1:
                            pose_ij.p = gymapi.Vec3(i*interval_xy, j*interval_xy, 1*interval_z)
                            pose_ij.r = corner_rot1
                            actor = self.gym.create_actor(env_, asset_blue_corner, pose_ij, "blue_corner" + str(i) + str(j), n, 0)
                        elif rot==2:
                            pose_ij.p = gymapi.Vec3(i*interval_xy, j*interval_xy, 1*interval_z)
                            pose_ij.r = corner_rot0
                            actor = self.gym.create_actor(env_, asset_blue_corner, pose_ij, "blue_corner" + str(i) + str(j), n, 0)
                        elif rot==3:
                            pose_ij.p = gymapi.Vec3(i*interval_xy, j*interval_xy, 1*interval_z)
                            pose_ij.r = corner_rot3
                            actor = self.gym.create_actor(env_, asset_blue_corner, pose_ij, "blue_corner" + str(i) + str(j), n, 0)
                        self.gym.set_actor_rigid_shape_properties(env_, actor, shape_props1)
                    elif tile_==7:
                        pose_ij.p = gymapi.Vec3(i*interval_xy, j*interval_xy, 1*interval_z)
                        actor = self.gym.create_actor(env_, asset_blue_cube, pose_ij, "blue_cube" + str(i) + str(j), n, 0)
                        if rot==0:
                            pose_ij.p = gymapi.Vec3(i*interval_xy, j*interval_xy, 2*interval_z)
                            pose_ij.r = corner_rot2
                            actor = self.gym.create_actor(env_, asset_yellow_corner, pose_ij, "yellow_corner" + str(i) + str(j), n, 0)
                        elif rot==1:
                            pose_ij.p = gymapi.Vec3(i*interval_xy, j*interval_xy, 2*interval_z)
                            pose_ij.r = corner_rot1
                            actor = self.gym.create_actor(env_, asset_yellow_corner, pose_ij, "yellow_corner" + str(i) + str(j), n, 0)
                        elif rot==2:
                            pose_ij.p = gymapi.Vec3(i*interval_xy, j*interval_xy, 2*interval_z)
                            pose_ij.r = corner_rot0
                            actor = self.gym.create_actor(env_, asset_yellow_corner, pose_ij, "yellow_corner" + str(i) + str(j), n, 0)
                        elif rot==3:
                            pose_ij.p = gymapi.Vec3(i*interval_xy, j*interval_xy, 2*interval_z)
                            pose_ij.r = corner_rot3
                            actor = self.gym.create_actor(env_, asset_yellow_corner, pose_ij, "yellow_corner" + str(i) + str(j), n, 0)
                        self.gym.set_actor_rigid_shape_properties(env_, actor, shape_props1)
                    elif tile_==8:
                        pose_ij.p = gymapi.Vec3(i*interval_xy, j*interval_xy, 1*interval_z)
                        actor = self.gym.create_actor(env_, asset_blue_cube, pose_ij, "blue_cube" + str(i) + str(j), n, 0)
                        pose_ij.p = gymapi.Vec3(i*interval_xy, j*interval_xy, 2*interval_z)
                        actor = self.gym.create_actor(env_, asset_yellow_cube, pose_ij, "yellow_cube" + str(i) + str(j), n, 0)
                        if rot==0:
                            pose_ij.p = gymapi.Vec3(i*interval_xy, j*interval_xy, 3*interval_z)
                            pose_ij.r = corner_rot2
                            actor = self.gym.create_actor(env_, asset_red_corner, pose_ij, "red_corner" + str(i) + str(j), n, 0)
                        elif rot==1:
                            pose_ij.p = gymapi.Vec3(i*interval_xy, j*interval_xy, 3*interval_z)
                            pose_ij.r = corner_rot1
                            actor = self.gym.create_actor(env_, asset_red_corner, pose_ij, "red_corner" + str(i) + str(j), n, 0)
                        elif rot==2:
                            pose_ij.p = gymapi.Vec3(i*interval_xy, j*interval_xy, 3*interval_z)
                            pose_ij.r = corner_rot0
                            actor = self.gym.create_actor(env_, asset_red_corner, pose_ij, "red_corner" + str(i) + str(j), n, 0)
                        elif rot==3:
                            pose_ij.p = gymapi.Vec3(i*interval_xy, j*interval_xy, 3*interval_z)
                            pose_ij.r = corner_rot3
                            actor = self.gym.create_actor(env_, asset_red_corner, pose_ij, "red_corner" + str(i) + str(j), n, 0)
                        self.gym.set_actor_rigid_shape_properties(env_, actor, shape_props1)
                    elif tile_==9:
                        pose_ij.p = gymapi.Vec3(i*interval_xy, j*interval_xy, 1*interval_z)
                        actor = self.gym.create_actor(env_, asset_blue_cube, pose_ij, "blue_cube" + str(i) + str(j), n, 0)
                        pose_ij.p = gymapi.Vec3(i*interval_xy, j*interval_xy, 2*interval_z)
                        actor = self.gym.create_actor(env_, asset_yellow_cube, pose_ij, "yellow_cube" + str(i) + str(j), n, 0)
                        pose_ij.p = gymapi.Vec3(i*interval_xy, j*interval_xy, 3*interval_z)
                        actor = self.gym.create_actor(env_, asset_red_cube, pose_ij, "red_cube" + str(i) + str(j), n, 0)
                        if rot==0:
                            pose_ij.p = gymapi.Vec3(i*interval_xy, j*interval_xy, 4*interval_z)
                            pose_ij.r = corner_rot2
                            actor = self.gym.create_actor(env_, asset_white_corner, pose_ij, "white_corner" + str(i) + str(j), n, 0)
                        elif rot==1:
                            pose_ij.p = gymapi.Vec3(i*interval_xy, j*interval_xy, 4*interval_z)
                            pose_ij.r = corner_rot1
                            actor = self.gym.create_actor(env_, asset_white_corner, pose_ij, "white_corner" + str(i) + str(j), n, 0)
                        elif rot==2:
                            pose_ij.p = gymapi.Vec3(i*interval_xy, j*interval_xy, 4*interval_z)
                            pose_ij.r = corner_rot0
                            actor = self.gym.create_actor(env_, asset_white_corner, pose_ij, "white_corner" + str(i) + str(j), n, 0)
                        elif rot==3:
                            pose_ij.p = gymapi.Vec3(i*interval_xy, j*interval_xy, 4*interval_z)
                            pose_ij.r = corner_rot3
                            actor = self.gym.create_actor(env_, asset_white_corner, pose_ij, "white_corner" + str(i) + str(j), n, 0)
                        self.gym.set_actor_rigid_shape_properties(env_, actor, shape_props1)
                    # ramps
                    elif tile_==10:
                        if rot==0:
                            pose_ij.p = gymapi.Vec3(i*interval_xy, j*interval_xy, 1*interval_z)
                            pose_ij.r = ramp_rot0
                            actor = self.gym.create_actor(env_, asset_blue_ramp, pose_ij, "blue_ramp" + str(i) + str(j), n, 0)
                        elif rot==1:
                            pose_ij.p = gymapi.Vec3(i*interval_xy, j*interval_xy, 1*interval_z)
                            pose_ij.r = ramp_rot1
                            actor = self.gym.create_actor(env_, asset_blue_ramp, pose_ij, "blue_ramp" + str(i) + str(j), n, 0)
                        elif rot==2:
                            pose_ij.p = gymapi.Vec3(i*interval_xy, j*interval_xy, 1*interval_z)
                            pose_ij.r = ramp_rot2
                            actor = self.gym.create_actor(env_, asset_blue_ramp, pose_ij, "blue_ramp" + str(i) + str(j), n, 0)
                        elif rot==3:
                            pose_ij.p = gymapi.Vec3(i*interval_xy, j*interval_xy, 1*interval_z)
                            pose_ij.r = ramp_rot3
                            actor = self.gym.create_actor(env_, asset_blue_ramp, pose_ij, "blue_ramp" + str(i) + str(j), n, 0)
                        self.gym.set_actor_rigid_shape_properties(env_, actor, shape_props)
                    elif tile_==11:
                        pose_ij.p = gymapi.Vec3(i*interval_xy, j*interval_xy, 1*interval_z)
                        actor = self.gym.create_actor(env_, asset_blue_cube, pose_ij, "blue_cube" + str(i) + str(j), n, 0)
                        if rot==0:
                            pose_ij.p = gymapi.Vec3(i*interval_xy, j*interval_xy, 2*interval_z)
                            pose_ij.r = ramp_rot0
                            actor = self.gym.create_actor(env_, asset_yellow_ramp, pose_ij, "yellow_ramp" + str(i) + str(j), n, 0)
                        elif rot==1:
                            pose_ij.p = gymapi.Vec3(i*interval_xy, j*interval_xy, 2*interval_z)
                            pose_ij.r = ramp_rot1
                            actor = self.gym.create_actor(env_, asset_yellow_ramp, pose_ij, "yellow_ramp" + str(i) + str(j), n, 0)
                        elif rot==2:
                            pose_ij.p = gymapi.Vec3(i*interval_xy, j*interval_xy, 2*interval_z)
                            pose_ij.r = ramp_rot2
                            actor = self.gym.create_actor(env_, asset_yellow_ramp, pose_ij, "yellow_ramp" + str(i) + str(j), n, 0)
                        elif rot==3:
                            pose_ij.p = gymapi.Vec3(i*interval_xy, j*interval_xy, 2*interval_z)
                            pose_ij.r = ramp_rot3
                            actor = self.gym.create_actor(env_, asset_yellow_ramp, pose_ij, "yellow_ramp" + str(i) + str(j), n, 0)
                        self.gym.set_actor_rigid_shape_properties(env_, actor, shape_props)
                    elif tile_==12:
                        pose_ij.p = gymapi.Vec3(i*interval_xy, j*interval_xy, 1*interval_z)
                        actor = self.gym.create_actor(env_, asset_blue_cube, pose_ij, "blue_cube" + str(i) + str(j), n, 0)
                        pose_ij.p = gymapi.Vec3(i*interval_xy, j*interval_xy, 2*interval_z)
                        actor = self.gym.create_actor(env_, asset_yellow_cube, pose_ij, "yellow_cube" + str(i) + str(j), n, 0)
                        if rot==0:
                            pose_ij.p = gymapi.Vec3(i*interval_xy, j*interval_xy, 3*interval_z)
                            pose_ij.r = ramp_rot0
                            actor = self.gym.create_actor(env_, asset_red_ramp, pose_ij, "red_ramp" + str(i) + str(j), n, 0)
                        elif rot==1:
                            pose_ij.p = gymapi.Vec3(i*interval_xy, j*interval_xy, 3*interval_z)
                            pose_ij.r = ramp_rot1
                            actor = self.gym.create_actor(env_, asset_red_ramp, pose_ij, "red_ramp" + str(i) + str(j), n, 0)
                        elif rot==2:
                            pose_ij.p = gymapi.Vec3(i*interval_xy, j*interval_xy, 3*interval_z)
                            pose_ij.r = ramp_rot2
                            actor = self.gym.create_actor(env_, asset_red_ramp, pose_ij, "red_ramp" + str(i) + str(j), n, 0)
                        elif rot==3:
                            pose_ij.p = gymapi.Vec3(i*interval_xy, j*interval_xy, 3*interval_z)
                            pose_ij.r = ramp_rot3
                            actor = self.gym.create_actor(env_, asset_red_ramp, pose_ij, "red_ramp" + str(i) + str(j), n, 0)
                        self.gym.set_actor_rigid_shape_properties(env_, actor, shape_props)
                    elif tile_==13:
                        pose_ij.p = gymapi.Vec3(i*interval_xy, j*interval_xy, 1*interval_z)
                        actor = self.gym.create_actor(env_, asset_blue_cube, pose_ij, "blue_cube" + str(i) + str(j), n, 0)
                        pose_ij.p = gymapi.Vec3(i*interval_xy, j*interval_xy, 2*interval_z)
                        actor = self.gym.create_actor(env_, asset_yellow_cube, pose_ij, "yellow_cube" + str(i) + str(j), n, 0)
                        pose_ij.p = gymapi.Vec3(i*interval_xy, j*interval_xy, 3*interval_z)
                        actor = self.gym.create_actor(env_, asset_red_cube, pose_ij, "red_cube" + str(i) + str(j), n, 0)
                        if rot==0:
                            pose_ij.p = gymapi.Vec3(i*interval_xy, j*interval_xy, 4*interval_z)
                            pose_ij.r = ramp_rot0
                            actor = self.gym.create_actor(env_, asset_white_ramp, pose_ij, "white_ramp" + str(i) + str(j), n, 0)
                        elif rot==1:
                            pose_ij.p = gymapi.Vec3(i*interval_xy, j*interval_xy, 4*interval_z)
                            pose_ij.r = ramp_rot1
                            actor = self.gym.create_actor(env_, asset_white_ramp, pose_ij, "white_ramp" + str(i) + str(j), n, 0)
                        elif rot==2:
                            pose_ij.p = gymapi.Vec3(i*interval_xy, j*interval_xy, 4*interval_z)
                            pose_ij.r = ramp_rot2
                            actor = self.gym.create_actor(env_, asset_white_ramp, pose_ij, "white_ramp" + str(i) + str(j), n, 0)
                        elif rot==3:
                            pose_ij.p = gymapi.Vec3(i*interval_xy, j*interval_xy, 4*interval_z)
                            pose_ij.r = ramp_rot3
                            actor = self.gym.create_actor(env_, asset_white_ramp, pose_ij, "white_ramp" + str(i) + str(j), n, 0)
                        self.gym.set_actor_rigid_shape_properties(env_, actor, shape_props)

            # generate random x,y coordinates in range [0,18] for the actor
            x = random.uniform(0,16)
            y = random.uniform(0,16)
            facing_ = 0.25 * (float)(random.randint(0,8))
            r_ = gymapi.Quat(0,1,0,1) * gymapi.Quat.from_axis_angle(gymapi.Vec3(1, 0, 0), facing_*math.pi)

            # create capsule_asset actor in the environment
            pose = gymapi.Transform()
            pose.p = gymapi.Vec3(x, y, self.initial_height)
            pose.r = gymapi.Quat(0,1,0,1)
            cap_handle = self.gym.create_actor(env_, capsule_asset, pose, 'capsule_asset', n, 0)
            self.gym.set_actor_rigid_shape_properties(env_, cap_handle, shape_props2)
            self.gym.set_actor_scale(env_, cap_handle, 0.4)
            self.actor_handles.append(cap_handle)
            self.gym.set_rigid_body_color(env_, cap_handle, 0, gymapi.MESH_VISUAL_AND_COLLISION, color_red)

            # names = self.gym.get_actor_rigid_body_names(env_, cap_handle)
            # print(names)
            
            # set random position for food
            x = random.uniform(0,16)
            y = random.uniform(0,16)
            # create capsule_asset food in the environment
            pose = gymapi.Transform()
            pose.p = gymapi.Vec3(x, y, self.initial_height)
            pose.r = gymapi.Quat(0,1,0,1)
            food_handle = self.gym.create_actor(env_, capsule_asset, pose, 'food_asset', n, 0)
            self.gym.set_actor_rigid_shape_properties(env_, food_handle, shape_props2)
            self.gym.set_actor_scale(env_, food_handle, 0.6)
            self.food_handles.append(food_handle)
            self.gym.set_rigid_body_color(env_, food_handle, 0, gymapi.MESH_VISUAL_AND_COLLISION, color_green)


            att_pose = gymapi.Transform()
            att_pose.p = gymapi.Vec3(x,y, self.initial_height)
            att_pose.r = r_
            attractor_properties_ = gymapi.AttractorProperties()
            attractor_properties_.stiffness = 5e10
            attractor_properties_.damping = 0
            attractor_properties_.axes = gymapi.AXIS_ROTATION
            attractor_properties_.target = att_pose
            attractor_properties_.rigid_handle = cap_handle
            attractor_handle_ = self.gym.create_rigid_body_attractor(env_, attractor_properties_)

            attractor_properties_food = gymapi.AttractorProperties()
            attractor_properties_food.stiffness = 5e10
            attractor_properties_food.damping = 0
            attractor_properties_food.axes = gymapi.AXIS_ROTATION
            attractor_properties_food.target = pose
            attractor_properties_food.rigid_handle = food_handle
            attractor_handle_food = self.gym.create_rigid_body_attractor(env_, attractor_properties_food)
            
            h1 = self.gym.create_camera_sensor(env_, camera_properties)
            camera_offset = gymapi.Vec3(-2, 0, 0)
            camera_rotation = gymapi.Quat.from_axis_angle(gymapi.Vec3(0, 0, 1), 0.4*math.pi)
            body_handle = self.gym.get_actor_rigid_body_handle(env_, cap_handle, 0)
            self.gym.attach_camera_to_body(h1, env_, body_handle, gymapi.Transform(camera_offset, camera_rotation), gymapi.FOLLOW_TRANSFORM)
            self.camera_handles.append(h1)
            self.attractor_handles.append(attractor_handle_)

            actor_count = self.gym.get_actor_count(env_)
            # # create blank actors
            # num_blk_act = max_actors_num_per_env - actor_count
            # for i in range(num_blk_act):
            #     pose = gymapi.Transform()
            #     pose.p = gymapi.Vec3(0, 0, 0)
            #     pose.r = gymapi.Quat(0,1,0,1)
            #     actor = self.gym.create_actor(env_, box_asset, pose, 'box_asset' + str(i), n, 0)
            #     self.gym.set_actor_scale(env_, actor, 0.000001)

        # a must-have operation 
        self.gym.prepare_sim(self.sim)

        self.step_count = 0

        if not self.headless:
            cam_pos = gymapi.Vec3(0.0, 15.0, 15.0)
            cam_target = gymapi.Vec3(1.5, 4, 0)
            self.gym.viewer_camera_look_at(self.viewer, None, cam_pos, cam_target)

        # time test
        self.start_time = time.time()
        self.last_frame_cnt = 0

        # Create helper geometry used for visualization
        # Create an wireframe axis
        self.axes_geom = gymutil.AxesGeometry(2)

        self.facing = 0
        self.facing_step = 0.25

        # local coordinate system
        self.verts = np.empty((3, 1), gymapi.Vec3.dtype)
        self.verts[0][0] = (1, 0, 0)
        self.verts[1][0] = (0, 1, 0)
        self.verts[2][0] = (0, 0, 1)

        # store initial status for all envs
        self.sim_status = []
        for i in range(self.num_matrix_envs):
            env_status = np.copy(self.gym.get_env_rigid_body_states(self.envs[i], gymapi.STATE_ALL))
            self.sim_status.append(env_status)
        
    @property
    def num_envs(self):
        return self.num_matrix_envs

    def set_world(self, seed_):

        pass

    def calculate_reward(self, env_, actor_handle, food_handle):

        i = 0
        body_states_actor = self.gym.get_actor_rigid_body_states(self.envs[i], self.actor_handles[i], gymapi.STATE_ALL)
        actor_p = body_states_actor['pose']['p']
        body_states_food = self.gym.get_actor_rigid_body_states(self.envs[i], self.food_handles[i], gymapi.STATE_ALL)
        food_p = body_states_food['pose']['p']

        # calculate distance between actor and food
        ap = np.array([actor_p[0][0], actor_p[0][1], actor_p[0][2]])
        fp = np.array([food_p[0][0], food_p[0][1], food_p[0][2]])
        dist = np.linalg.norm(ap - fp)
        # print(dist)

        if dist < self.proximity_threshold:
            # print('reward: 1')
            return 1
        else:
            return 0

    def done_determinant(self, env_, actor_handle, food_handle, reward = 0, step = 0):

        if reward == 1:
            return True

        i = 0
        body_states_actor = self.gym.get_actor_rigid_body_states(self.envs[i], self.actor_handles[i], gymapi.STATE_ALL)
        actor_p = body_states_actor['pose']['p']
        actor_cnnc_x = ((actor_p[0][0]+1)/(9*2))*20*9
        actor_cnnc_y = ((actor_p[0][1]+1)/(9*2))*20*9

        body_states_food = self.gym.get_actor_rigid_body_states(self.envs[i], self.food_handles[i], gymapi.STATE_ALL)
        food_p = body_states_food['pose']['p']
        food_cnnc_x = ((food_p[0][0]+1)/(9*2))*20*9
        food_cnnc_y = ((food_p[0][1]+1)/(9*2))*20*9

        if not (actor_cnnc_x >= 0 and actor_cnnc_x < 20*9 and actor_cnnc_y >= 0 and actor_cnnc_y < 20*9):
            return True
        
        if not (food_cnnc_x >= 0 and food_cnnc_x < 20*9 and food_cnnc_y >= 0 and food_cnnc_y < 20*9):
            return True

        actor_cnnt_idx = self.connectivity_map[int(actor_cnnc_y)][int(actor_cnnc_x)]
        food_cnnt_idx = self.connectivity_map[int(food_cnnc_y)][int(food_cnnc_x)]

        if actor_cnnt_idx != 1 and food_cnnt_idx != 1:
            if actor_cnnt_idx != food_cnnt_idx:
                return True

        if step > self.max_steps:
            return True

        return False

    def step(self, action : int, physics_only = False):

        self.step_count += 1

        l_ = int((HEIGHT-1)/5)

        i = 0

        '''
        apply action
        '''
        body_states = self.gym.get_actor_rigid_body_states(self.envs[i], self.actor_handles[i], gymapi.STATE_ALL)
        # transform global coordinate system to local coordinate system
        body_ = self.gym.get_actor_rigid_body_handle(self.envs[i], self.actor_handles[i], 0)
        body_t = self.gym.get_rigid_transform(self.envs[i], body_)
        body_t.p = gymapi.Vec3(0,0,0)
        verts_ = body_t.transform_points(self.verts)

        # maintain vertical speed
        z_spd = body_states['vel']['linear'][0][2]

        if action == 0:
            body_states['vel']['linear'].fill((self.spd * verts_[1][0][0],self.spd * verts_[1][0][1],z_spd))
            self.gym.set_actor_rigid_body_states(self.envs[i], self.actor_handles[i], body_states, gymapi.STATE_VEL)
        elif action == 1:
            body_states['vel']['linear'].fill((-self.spd * verts_[1][0][0],-self.spd * verts_[1][0][1],z_spd))
            self.gym.set_actor_rigid_body_states(self.envs[i], self.actor_handles[i], body_states, gymapi.STATE_VEL)
        elif action == 2:
            body_states['vel']['linear'].fill((-self.spd * verts_[2][0][0],-self.spd * verts_[2][0][1],z_spd))
            self.gym.set_actor_rigid_body_states(self.envs[i], self.actor_handles[i], body_states, gymapi.STATE_VEL)
        elif action == 3:
            body_states['vel']['linear'].fill((self.spd * verts_[2][0][0],self.spd * verts_[2][0][1],z_spd))
            self.gym.set_actor_rigid_body_states(self.envs[i], self.actor_handles[i], body_states, gymapi.STATE_VEL)
        elif action == 4:
            self.facing -= self.facing_step
            attractor_pose = gymapi.Transform()
            attractor_pose.p = gymapi.Vec3(0, 0.0, 0.0)
            attractor_pose.r = gymapi.Quat(0,1,0,1) * gymapi.Quat.from_axis_angle(gymapi.Vec3(1, 0, 0), self.facing*math.pi)
            self.gym.set_attractor_target(self.envs[i], self.attractor_handles[i], attractor_pose)
        elif action == 5:
            self.facing += self.facing_step
            attractor_pose = gymapi.Transform()
            attractor_pose.p = gymapi.Vec3(0, 0.0, 0.0)
            attractor_pose.r = gymapi.Quat(0,1,0,1) * gymapi.Quat.from_axis_angle(gymapi.Vec3(1, 0, 0), self.facing*math.pi)
            self.gym.set_attractor_target(self.envs[i], self.attractor_handles[i], attractor_pose)
        # # break !!
        # elif action == 6:
        #     body_states['vel']['linear'].fill((0,0,0))
        #     self.gym.set_actor_rigid_body_states(self.envs[i], self.actor_handles[i], body_states, gymapi.STATE_VEL)

        observation = np.zeros((HEIGHT,WIDTH, 4), np.uint8)

        '''
        step the physics
        '''
        self.gym.simulate(self.sim)
        self.gym.fetch_results(self.sim, True)
        if not self.headless:
            self.gym.clear_lines(self.viewer)
        self.gym.step_graphics(self.sim)

        if not physics_only:

            self.gym.render_all_camera_sensors(self.sim)

            '''
            get observation
            '''
            rgb_image = self.gym.get_camera_image(self.sim, self.envs[i], self.camera_handles[i], gymapi.IMAGE_COLOR)
            # create an opencv image of size (HEIGHT,WIDTH , 4)
            
            observation[:, :, :4] = rgb_image.reshape((HEIGHT,WIDTH, 4))
            # rotate 90 degrees counter-clockwise
            observation = np.rot90(observation, -1)
            # transform the image from RGBA to RGB
            observation = observation[:, :, :3]
            # transpose the image from (HEIGHT,WIDTH, 3) to (3, HEIGHT, WIDTH)
            observation = np.transpose(observation, (2, 0, 1))

            # if self.render_indicator:
            #     # velocity indicator
            #     body_ti = body_t.inverse()
            #     vel_ = np.empty((1, 1), gymapi.Vec3.dtype)
            #     vel_[0][0] = (body_states['vel']['linear'][0][0], body_states['vel']['linear'][0][1], body_states['vel']['linear'][0][2])
            #     vel_ego = body_ti.transform_points(vel_)
            #     # print(vel_ego)
            #     if vel_ego[0][0][1]>=0 :
            #         h_ = l_*(vel_ego[0][0][1]/self.spd)
            #         h_ = min(l_,h_)
            #         observation[:12,:(int)(h_),:]=(200,200,200)
            #     elif vel_ego[0][0][1]<0 :
            #         h_ = -l_*(vel_ego[0][0][1]/self.spd)
            #         h_ = min(l_,h_)
            #         observation[:12,l_:l_+(int)(h_),:]=(200,200,200)
            #     if vel_ego[0][0][2]>=0 :
            #         h_ = l_*(vel_ego[0][0][2]/self.spd)
            #         h_ = min(l_,h_)
            #         observation[:12,l_*2:l_*2+(int)(h_),:]=(200,200,200)
            #     elif vel_ego[0][0][2]<0 :
            #         h_ = -l_*(vel_ego[0][0][2]/self.spd)
            #         h_ = min(l_,h_)
            #         observation[:12,l_*3:l_*3+(int)(h_),:]=(200,200,200)
            #     # orientation indicator
            #     dir_ = abs(self.facing) % 2
            #     h_ = l_*(dir_/2)
            #     h_ = min(l_,h_)
            #     observation[:12,l_*4:l_*4+(int)(h_),:]=(200,200,200)

            self.camera_sensor_image = np.copy(np.transpose(observation, (1, 2, 0)))

        else:
            observation = np.zeros((N_CHANNELS, HEIGHT,WIDTH), np.uint8)

        '''
        get reward
        '''
        reward = self.calculate_reward(self.envs[i], self.actor_handles[i], self.food_handles[i])
        done = self.done_determinant(self.envs[i], self.actor_handles[i], self.food_handles[i],reward,self.step_count)
        info = {}

        if self.show_action:
            self.show_image()
        
        # if not self.headless:
        #     # update the viewer
        #     self.gym.draw_viewer(self.viewer, self.sim, True)

        # end_time = time.time()
        # if (end_time - self.start_time) > 1:
        #     print("FPS: %.2f" % ((self.step_count - self.last_frame_cnt)*self.num_matrix_envs))
        #     self.last_frame_cnt = self.step_count
        #     self.start_time = time.time()
        
        return np.copy(observation), reward, done, info

    def reset(self):

        self.step_count = 0
        self.facing = 0
        self.start_time = time.time()

        i = 0

        done_at_begining = False #self.done_determinant(self.envs[i], self.actor_handles[i], self.food_handles[i])

        self.gym.set_env_rigid_body_states(self.envs[i], self.sim_status[i], gymapi.STATE_ALL)

        # reset random position and orientation for actor
        x = random.uniform(0,16)
        y = random.uniform(0,16)

        reset_height = self.initial_height

        facing_i = (float)(random.randint(0,8))
        facing_ = 0.25 * facing_i
        self.facing = facing_
        r_ = gymapi.Quat(0,1,0,1) * gymapi.Quat.from_axis_angle(gymapi.Vec3(1, 0, 0), facing_*math.pi)
        body_states = self.gym.get_actor_rigid_body_states(self.envs[i], self.actor_handles[i], gymapi.STATE_ALL)
        body_states['pose']['p'].fill((x,y,reset_height))
        self.gym.set_actor_rigid_body_states(self.envs[i], self.actor_handles[i], body_states, gymapi.STATE_POS)
        attractor_pose = gymapi.Transform()
        attractor_pose.p = gymapi.Vec3(0, 0.0, 0.0)
        attractor_pose.r = r_
        self.gym.set_attractor_target(self.envs[i], self.attractor_handles[i], attractor_pose)
        # print("reset actor at (%d, %d)" % (x, y))
        # reset random position for food
        x = random.uniform(0,16)
        y = random.uniform(0,16)
        body_states = self.gym.get_actor_rigid_body_states(self.envs[i], self.food_handles[i], gymapi.STATE_ALL)
        body_states['pose']['p'].fill((x,y,reset_height))
        self.gym.set_actor_rigid_body_states(self.envs[i], self.food_handles[i], body_states, gymapi.STATE_POS)

        # reset actor vel=0
        body_states = self.gym.get_actor_rigid_body_states(self.envs[i], self.actor_handles[i], gymapi.STATE_ALL)
        body_states['vel']['linear'].fill((0,0,0))
        self.gym.set_actor_rigid_body_states(self.envs[i], self.actor_handles[i], body_states, gymapi.STATE_VEL)

        # step the physics
        self.gym.simulate(self.sim)
        self.gym.fetch_results(self.sim, True)
        self.gym.step_graphics(self.sim)
        self.gym.render_all_camera_sensors(self.sim)

        rgb_image = self.gym.get_camera_image(self.sim, self.envs[i], self.camera_handles[i], gymapi.IMAGE_COLOR)
        # create an opencv image of size (HEIGHT,WIDTH , 4)
        observation = np.zeros((HEIGHT,WIDTH, 4), np.uint8)
        observation[:, :, :4] = rgb_image.reshape((HEIGHT,WIDTH, 4))
        # rotate 90 degrees counter-clockwise
        observation = np.rot90(observation, -1)
        # transform the image from RGBA to RGB
        observation = observation[:, :, :3]
        # transpose the image from (HEIGHT,WIDTH, 3) to (3, HEIGHT, WIDTH)
        observation = np.transpose(observation, (2, 0, 1))
        # print("shape of observation: ", observation.shape)

        # '''
        # simulate until actor and food lands on the terrain
        # '''
        # for t in range(400):
        #     observation, reward, done, info = self.step(action = -3, physics_only = True)

        # # done_at_begining = not self.done_determinant(self.envs[i], self.actor_handles[i], self.food_handles[i])
        
        # observation, reward, done, info = self.step(action = -3, physics_only = False)
        # transpose the image from (3, HEIGHT, WIDTH) to (HEIGHT,WIDTH, 3)
        self.camera_sensor_image = np.copy(np.transpose(observation, (1, 2, 0)))

        return np.copy(observation)

    def render(self, mode='human', real_time = False):
        
        i = 0
        body_states = self.gym.get_actor_rigid_body_states(self.envs[i], self.actor_handles[i], gymapi.STATE_ALL)
        actor_p = body_states['pose']['p']
        # print("actor position: (%f, %f, %f)" % (actor_p[0][0], actor_p[0][1], actor_p[0][2]))
        cnnc_x = ((actor_p[0][0]+1)/(9*2))*20*9
        cnnc_y = ((actor_p[0][1]+1)/(9*2))*20*9

        # print(self.connectivity_map[int(cnnc_y)][int(cnnc_x)])

        if not self.headless:

            if self.gym.query_viewer_has_closed(self.viewer):
                return -2

            # for i in range(self.num_matrix_envs):
            body_states = self.gym.get_camera_transform(self.sim, self.envs[i], self.camera_handles[i])
            draw_camera(self.gym, self.viewer,self.envs[i],body_states)
            # self.gym.draw_env_rigid_contacts(self.viewer, self.envs[i], gymapi.Vec3(1,0,0), 2, True)
            # update the viewer
            self.gym.draw_viewer(self.viewer, self.sim, True)

            if real_time:
                self.gym.sync_frame_time(self.sim)

            self.camera_sensor_image = cv2.cvtColor(self.camera_sensor_image, cv2.COLOR_RGB2BGR)
            cv2.imshow('rgb', self.camera_sensor_image)

            # draw dot at [cnnc_x, cnnc_y]
            cnnc_vis = np.copy(self.connectivity_map_visual)
            cv2.circle(cnnc_vis, (int(cnnc_x), int(cnnc_y)), 5, (0, 0, 255), -1)
            cv2.imshow('conn',cnnc_vis)
            self.interaction = cv2.waitKey(1)
            
        return self.interaction

    def show_image(self):
        
        i = 0
        body_states = self.gym.get_actor_rigid_body_states(self.envs[i], self.actor_handles[i], gymapi.STATE_ALL)
        actor_p = body_states['pose']['p']
        # print("actor position: (%f, %f, %f)" % (actor_p[0][0], actor_p[0][1], actor_p[0][2]))
        cnnc_x = ((actor_p[0][0]+1)/(9*2))*20*9
        cnnc_y = ((actor_p[0][1]+1)/(9*2))*20*9

        self.camera_sensor_image = cv2.cvtColor(self.camera_sensor_image, cv2.COLOR_RGB2BGR)
        cv2.imshow('rgb', self.camera_sensor_image)
        cv2.waitKey(1)

    def close (self):
        
        if not self.headless:
            self.gym.destroy_viewer(self.viewer)

        self.gym.destroy_sim(self.sim)

    def __enter__(self):
        """Support with-statement for the environment."""
        return self

    def __exit__(self, *args):
        """Support with-statement for the environment."""
        self.close()
        # propagate exception
        return False